package com.wjd.store;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Stream;

/**
 * 记录树存储结构
 *
 * @author weijiaduo
 * @since 2024/3/13
 */
public abstract class RecordTree<T extends Record> {

    /**
     * 读写锁
     */
    private static final ReadWriteLock LOCK = new ReentrantReadWriteLock();

    /**
     * 词语 id 文件名
     */
    private final String ID_FILE = "ids";
    /**
     * id 对应的数据文件目录
     */
    private final String ID_DATA_DIR = "id_data";
    /**
     * 数据文件扩展名
     */
    private final String DATA_EXT;
    /**
     * 根目录路径
     */
    private final Path basePath;
    /**
     * id 文件
     */
    private final Path idPath;
    /**
     * 数据目录路径
     */
    private final Path dataPath;
    /**
     * 缓存已有的词语及其 id
     */
    private final HashMap<String, Long> key2Id = new HashMap<>();
    /**
     * 当前最大的 id
     */
    private long maxKeyId = 0;
    /**
     * 树高度
     */
    private int level;

    public RecordTree(String baseDir) throws IOException {
        this(baseDir, ".fh");
    }

    public RecordTree(String baseDir, String ext) throws IOException {
        this.DATA_EXT = ext;

        // 先创建必要的文件和目录
        // 根目录
        basePath = Paths.get(baseDir);
        if (!Files.exists(basePath)) {
            Files.createDirectories(basePath);
        }
        // id 文件
        idPath = basePath.resolve(ID_FILE);
        if (!Files.exists(idPath)) {
            Files.createFile(idPath);
        }
        // id 数据目录
        dataPath = basePath.resolve(ID_DATA_DIR);
        if (!Files.exists(dataPath)) {
            Files.createDirectories(dataPath);
        }

        // 以 0 开头的目录/文件，表示当前层级
        File baseFile = basePath.toFile();
        File[] files = baseFile.listFiles((dir, name) -> name.startsWith("0"));
        if (files == null || files.length == 0) {
            level = 0;
        } else {
            String str = files[0].getName().substring(1);
            level = Integer.parseInt(str);
        }

        // 初始化 id
        initKeyIds();
    }

    /**
     * 从本地加载 id
     */
    private void initKeyIds() {
        try (FileInputStream fis = new FileInputStream(idPath.toFile());
             BufferedInputStream bis = new BufferedInputStream(fis);
             DataInputStream in = new DataInputStream(bis)) {
            key2Id.clear();
            maxKeyId = 0;
            while (in.available() != 0) {
                String url = in.readUTF();
                long id = in.readLong();
                key2Id.put(url, id);
                maxKeyId = Math.max(id, maxKeyId);
            }
        } catch (IOException e) {
            throw new RuntimeException("初始化数据库失败", e);
        }
    }

    /**
     * 查询指定 key 的数据记录
     *
     * @param key 指定 key
     * @return 记录集合
     */
    public List<T> queryByKey(String key) {
        Long id = key2Id.get(key);
        if (id == null) {
            return Collections.emptyList();
        }
        return queryById(id);
    }

    /**
     * 查询指定 id 的数据记录
     *
     * @param id 指定 id
     * @return 记录集合
     */
    public List<T> queryById(long id) {
        File file = findDataFile(id);
        if (!file.exists()) {
            return Collections.emptyList();
        }

        LOCK.readLock().lock();
        try (FileInputStream fis = new FileInputStream(file);
             BufferedInputStream bis = new BufferedInputStream(fis);
             DataInputStream in = new DataInputStream(bis)) {
            List<T> res = new ArrayList<>();
            while (in.available() != 0) {
                T record = readRecord(in);
                res.add(record);
            }
            return res;
        } catch (IOException e) {
            throw new RuntimeException("读取文件" + file + "失败！", e);
        } finally {
            LOCK.readLock().unlock();
        }
    }

    /**
     * 从输入流中读取一条记录返回
     *
     * @param in 输入流
     * @return 记录
     */
    protected abstract T readRecord(DataInputStream in) throws IOException;

    /**
     * 保存记录
     *
     * @param record 记录
     * @return 记录 id
     */
    public long save(T record) {
        // 添加 id
        long id = addKey(record.getKey());
        // 添加数据
        addRecord(id, record);
        return id;
    }

    /**
     * 当 key 不在数据库中，则插入 key 到数据库中
     *
     * @param key 词语
     * @return 词语对应的 id
     */
    private long addKey(String key) {
        LOCK.writeLock().lock();
        // 验证词语是否已存在
        Long id = key2Id.get(key);
        if (id != null) {
            return id;
        }
        try (FileOutputStream fos = new FileOutputStream(getIdFile(), true);
             BufferedOutputStream bos = new BufferedOutputStream(fos);
             DataOutputStream out = new DataOutputStream(bos)) {
            // 给词语一个 id
            id = ++maxKeyId;
            key2Id.put(key, id);

            // 保存词语和 id
            out.writeUTF(key);
            out.writeLong(id);
            return id;
        } catch (IOException e) {
            throw new RuntimeException("插入 key：" + key + "时失败", e);
        } finally {
            LOCK.writeLock().unlock();
        }
    }

    /**
     * 添加新纪录到指定文件（该文件存在）
     */
    private void addRecord(long id, T record) {
        LOCK.writeLock().lock();
        try (FileOutputStream fos = new FileOutputStream(getDataFile(id), true);
             BufferedOutputStream bos = new BufferedOutputStream(fos);
             DataOutputStream out = new DataOutputStream(bos)) {
            writeRecord(out, record);
        } catch (FileNotFoundException e) {
            throw new RuntimeException("未能找到文件" + id, e);
        } catch (Exception e) {
            throw new RuntimeException("写入文件" + id + "失败", e);
        } finally {
            LOCK.writeLock().unlock();
        }
    }

    /**
     * 将一条记录写入输出流中
     *
     * @param out    输出流
     * @param record 记录
     */
    protected abstract void writeRecord(DataOutputStream out, T record) throws IOException;

    /**
     * 获取记录词语 id 的文件
     * <p>
     * 注意：如果文件不存在，则会新建文件
     *
     * @return 保存词语 id 的文件
     */
    private File getIdFile() {
        return idPath.toFile();
    }

    /**
     * 获取指定 id 对应的数据文件
     * <p>
     * 注意：如果文件不存在，则会新建文件
     *
     * @param id 指定 id
     * @return id 对应的数据文件
     */
    private File getDataFile(long id) {
        try {
            IDPath idPath = new IDPath(id, level);
            // 确保文件存在
            ensurePath(idPath);
            return getPathFile(idPath);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 搜索指定 id 对应的文件
     * <p>
     * 注意：文件可能不存在，也不会新建文件
     *
     * @param id 指定 id
     * @return id 对应的文件
     */
    private File findDataFile(long id) {
        IDPath idPath = new IDPath(id, level);
        return getPathFile(idPath);
    }

    /**
     * 确保指定 id 对应的文件存在
     *
     * @param idPath id 路径
     */
    private void ensurePath(IDPath idPath) throws IOException {
        // 寻找文件前，先验证文件树高度
        String[] path = idPath.paths;
        if (level < path.length) {
            // 文件树高度不足，需要提升
            levelUp();
        }

        File file = getPathFile(idPath);
        if (!file.exists()) {
            // 文件不存在，则创建
            createPathFile(idPath);
        }
    }

    /**
     * 使文件树升高一层
     */
    private void levelUp() throws IOException {
        LOCK.writeLock().lock();
        try {
            Path dir = dataPath;
            // 提升树的高度
            int lev = ++level;

            // 第1层，只需要一个标记文件
            if (lev == 1) {
                Path file = dir.resolve("0" + lev + DATA_EXT);
                if (!Files.exists(file)) {
                    Files.createFile(file);
                }
                return;
            }

            // 移动文件到新目录中
            try (Stream<Path> ps = Files.list(dir)) {
                // 获取需要移动的文件
                List<Path> files = ps.toList();
                // 再创建目标子目录
                Path sub = dir.resolve("0" + lev);
                if (!Files.exists(sub)) {
                    Files.createDirectory(sub);
                }
                // 把根目录下的文件，移动到新子目录下
                for (Path file : files) {
                    Path newFile = sub.resolve(file.getFileName());
                    Files.move(file, newFile);
                }
            }
        } finally {
            LOCK.writeLock().unlock();
        }
    }

    /**
     * 在指定目录下创建文件路径
     */
    private void createPathFile(IDPath idPath) throws IOException {
        String[] path = idPath.paths;
        int length = path.length;
        if (length < 1) {
            return;
        }

        LOCK.writeLock().lock();
        try {
            // 创建目录
            Path dir = dataPath;
            for (int i = 0; i < length - 1; i++) {
                dir = dir.resolve(path[i]);
                if (!Files.exists(dir)) {
                    Files.createDirectory(dir);
                }
            }
            // 创建文件
            Path file = dir.resolve(path[length - 1] + DATA_EXT);
            if (!Files.exists(file)) {
                Files.createFile(file);
            }
        } finally {
            LOCK.writeLock().unlock();
        }
    }

    /**
     * 获取指定路径的文件
     *
     * @param idPath 指定路径
     * @return 文件
     */
    private File getPathFile(IDPath idPath) {
        Path dir = dataPath;
        String[] path = idPath.paths;
        int length = path.length;
        for (int i = 0; i < length - 1; i++) {
            dir = dir.resolve(path[i]);
        }
        Path file = dir.resolve(path[length - 1] + DATA_EXT);
        return file.toFile();
    }

    /**
     * 用于划分id变成路径的内部类
     */
    static class IDPath {

        // id 路径
        final String[] paths;

        public IDPath(long id, int level) {
            String[] ids = divide(id);
            paths = formatPath(ids, level);
        }

        /**
         * id 拆分成小段
         *
         * @param id id
         * @return 小段路径
         */
        private String[] divide(long id) {
            if (id == 0) {
                return new String[]{"0"};
            }
            int cnt = 0;
            for (long d = id; d > 0; d /= 100) {
                cnt++;
            }
            String[] ids = new String[cnt];
            long d = id;
            for (int i = cnt - 1; i >= 0; i--) {
                ids[i] = String.valueOf(d % 100);
                d /= 100;
            }
            return ids;
        }

        /**
         * 格式化路径
         *
         * @param ids 路径 id
         * @return 格式化路径
         */
        private String[] formatPath(String[] ids, int level) {
            int ol = ids.length;
            int nl = Math.max(ol, level);
            String[] path = new String[nl];
            for (int i = nl - 1, j = ol - 1; i >= 0; i--, j--) {
                String p = j >= 0 ? ids[j] : "0";
                if ("0".equals(p)) {
                    path[i] = "0" + (nl - i);
                } else {
                    path[i] = p;
                }
            }
            return path;
        }
    }

}
